import os
import yaml

# 项目目录结构
project_structure = {
    'dataset': ['类别1', '类别2', '类别3'],
    'logs': [],
    'models': [],
    'src': []
}

# 配置文件内容
config_content = {
    'data': {
        'data_dir': 'dataset',
        'img_size': 224,
        'batch_size': 32
    },
    'model': {
        'type': 'transfer',
        'save_dir': 'models',
        'learning_rate': 0.001
    },
    'training': {
        'epochs': 10,
        'validation_split': 0.2,
        'early_stopping_patience': 5
    }
}

def create_project_structure():
    """创建项目目录结构"""
    print("开始创建项目结构...")
    
    # 创建目录
    for directory, subdirs in project_structure.items():
        os.makedirs(directory, exist_ok=True)
        print(f"创建目录: {directory}")
        
        # 创建子目录
        for subdir in subdirs:
            path = os.path.join(directory, subdir)
            os.makedirs(path, exist_ok=True)
            print(f"创建子目录: {path}")

def create_config_file():
    """创建配置文件"""
    print("\n创建配置��件...")
    
    config_path = 'config.yml'
    with open(config_path, 'w', encoding='utf-8') as f:
        yaml.dump(config_content, f, default_flow_style=False, allow_unicode=True)
    print(f"配置文件已创建: {config_path}")

def main():
    """主函数"""
    print("开始设置项目...")
    
    # 创建项目结构
    create_project_structure()
    
    # 创建配置文件
    create_config_file()
    
    print("\n项目设置完成！")
    print("\n接下来你需要:")
    print("1. 将训练图片放入 dataset/类别X 目录")
    print("2. 根据需要修改 config.yml 配置文件")
    print("3. 运行 python src/main.py 开始训练")

if __name__ == "__main__":
    main() 